import numpy as np
import matplotlib.pyplot as plt
import math
import mpl_toolkits.axisartist as axisartist
from scipy import integrate


def square(t, T, delta):
    # 方波信号
    if (t + delta) % T <= 2 * delta:
        return 1
    else:
        return 0

def calculate_args(T, delta, func, n):
    def fun(t):
        if (t + delta) % T <= 2 * delta:
            return 1 * func(t, n, T)
        else:
            return 0
    tmp = 1 if n == 0 else 2
    return tmp / T * integrate.quad(fun, -T/2, T/2)[0]

def fcos(t, n, T):
    return math.cos(2 * math.pi * n * 1 / T * t)

def fsin(t, n, T):
    return math.sin(2 * math.pi * n * 1 / T * t)

def sum(n, T, delta, t):
    fun = np.zeros_like(t)
    for i in range(n):
        tmp = np.array([fcos(j, i, T) * calculate_args(T, delta, fcos, i) for j in t])
        fun += tmp
        if i != 0:
            tmp = np.array([fsin(j, i, T) * calculate_args(T, delta, fsin, i) for j in t])
            fun += tmp
    return fun

def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)


if __name__ == "__main__":

    fig = plt.figure(figsize=(12, 8))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    T = 8
    delta = 2
    t = np.arange(-2 * T, 2 * T, 0.01)

    fun1 = np.array([square(i, T, delta) for i in t])
    set_axis(fig, 221, "原始的方波信号", (-2 * T, 2 * T), (-1, 2))
    plt.plot(t, fun1)

    n = 5
    fun2 = sum(n, T, delta, t)
    set_axis(fig, 222, "n = 5时傅立叶反变换得到的图像", (-2 * T, 2 * T), (-1, 2))
    plt.plot(t, fun2)

    n = 20
    fun3 = sum(n, T, delta, t)
    set_axis(fig, 223, "n = 20时傅立叶反变换得到的图像", (-2 * T, 2 * T), (-1, 2))
    plt.plot(t, fun3)

    n = 100
    fun4 = sum(n, T, delta, t)
    set_axis(fig, 224, "n = 100时傅立叶反变换得到的图像", (-2 * T, 2 * T), (-1, 2))
    plt.plot(t, fun4)

    plt.show()